import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import glob

# 图像尺寸
imageWidth = 1280
imageHeight = 1024
imageSize = (imageWidth, imageHeight)

# 相机内参矩阵与外部参数
mtx_l = np.array([[1.933107611281841e+04, 0.00000000e+00, 9.694984440499944e+02, 0],
                  [0.00000000e+00, 1.934122507949005e+04, 2.956718768245792e+02, 0],
                  [0.00000000e+00, 0.00000000e+00, 1.00000000e+00, 0]])
mtx_r = np.array([[1.917063604337356e+04, 0.00000000e+00, 6.175710227944095e+02, 0],
                  [0.00000000e+00, 1.917584413947483e+04, 4.751439387234215e+02, 0],
                  [0.00000000e+00, 0.00000000e+00, 1.00000000e+00, 0]])
R_lr = np.array([[0.997507529504717, -9.289799944656619e-04, 0.070553990514821],
                 [0.001414684431041, 0.999975643865923, -0.006834496535399],
                 [-0.070545922981815, 0.006917273386363, 0.997484548291121]])
T_lr = np.array([[-8.893666081523952e+02],
                 [-25.589864825494843],
                 [-69.797697709734280]])
cameraMatrixL = mtx_l[:, 0:3]
cameraMatrixR = mtx_r[:, 0:3]

# 畸变参数
distCoeffL = np.array([-0.512681529616576, 1.706376412853077e+02, 6.775499882882220e-05,
                       -0.005634824339793, -3.531674034274161e+04])

distCoeffR = np.array([0.153626897215147, 1.542117056873553e+02, -0.001036188250721,
                       -0.002730184943375, 0.208635129241957])

# 左相机到左相机的投影矩阵
R_ll = ([[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]])
T_ll = ([[0], [0], [0]])
temp_R_ll = np.append(R_ll, T_ll, axis=1)
_temp_R_ll = np.row_stack((temp_R_ll, [0, 0, 0, 1]))
P0 = np.dot(mtx_l, _temp_R_ll)

# 左相机到右相机的投影矩阵
temp_R_lr = np.append(R_lr, T_lr, axis=1)
_temp_R_lr = np.row_stack((temp_R_lr, [0, 0, 0, 1]))
P1 = np.dot(mtx_r, _temp_R_lr)

# 计算立体校正参数
Rl, Rr, Pl, Pr, Q, validROIL, validROIR = cv2.stereoRectify(cameraMatrixL, distCoeffL,
                                                            cameraMatrixR, distCoeffR,
                                                            imageSize, R_lr, T_lr,
                                                            flags=0,
                                                            alpha=0, newImageSize=(0, 0))

# 计算更正map
mapLx, mapLy = cv2.initUndistortRectifyMap(cameraMatrixL, distCoeffL, Rl, Pl, imageSize, cv2.CV_32FC1)
mapRx, mapRy = cv2.initUndistortRectifyMap(cameraMatrixR, distCoeffR, Rr, Pr, imageSize, cv2.CV_32FC1)

# 读取图片
grayImageL = cv2.imread('L2.bmp', 0)
grayImageR = cv2.imread('R2.bmp', 0)

# 经过remap之后，左右相机的图像已经共面并且行对齐
rectifyImageL = cv2.remap(grayImageL, mapLx, mapLy, cv2.INTER_LINEAR)
rectifyImageR = cv2.remap(grayImageR, mapRx, mapRy, cv2.INTER_LINEAR)

# 立体矫正对比图
plt.figure(0)
plt.subplot(2, 2, 1)
plt.imshow(grayImageL, cmap='gray')
plt.title('ImageL')
plt.subplot(2, 2, 3)
plt.imshow(rectifyImageL, cmap='gray')
plt.title('rectifyImageL')

plt.subplot(2, 2, 2)
plt.imshow(grayImageR, cmap='gray')
plt.title('ImageR')
plt.subplot(2, 2, 4)
plt.imshow(rectifyImageR, cmap='gray')
plt.title('rectifyImageR')
plt.show()

# 设置棋盘格标定板行列数
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
retL, cornersL = cv2.findChessboardCornersSB(rectifyImageL, (15, 11), None)
retR, cornersR = cv2.findChessboardCornersSB(rectifyImageR, (15, 11), None)
cornersL = cv2.cornerSubPix(rectifyImageL, cornersL, (15, 11), (-1, -1), criteria)
cornersR = cv2.cornerSubPix(rectifyImageR, cornersR, (15, 11), (-1, -1), criteria)
for i in range(0, len(cornersL)):
    cv2.circle(rectifyImageL, (int(cornersL[i][0][0]), int(cornersL[i][0][1])), 1, (255, 0, 0))

retLO, cornersLO = cv2.findChessboardCornersSB(grayImageL, (15, 11), None)
retRO, cornersRO = cv2.findChessboardCornersSB(grayImageR, (15, 11), None)
cornersLO = cv2.cornerSubPix(grayImageL, cornersLO, (15, 11), (-1, -1), criteria)
cornersRO = cv2.cornerSubPix(grayImageR, cornersRO, (15, 11), (-1, -1), criteria)

# 绘制识别到的棋盘格角点
rectifyImageR = rectifyImageR[:, :, np.newaxis]
rectifyImageRR = rectifyImageR.repeat([3], axis=2)
for i in range(0, len(cornersR)):
    cv2.line(rectifyImageRR, (int(cornersR[i][0][0]) - 3, int(cornersR[i][0][1])),
             (int(cornersR[i][0][0]) + 3, int(cornersR[i][0][1])), color=(0, 0, 255), thickness=1)
    cv2.line(rectifyImageRR, (int(cornersR[i][0][0]), int(cornersR[i][0][1]) - 3),
             (int(cornersR[i][0][0]), int(cornersR[i][0][1]) + 3), color=(0, 0, 255), thickness=1)

cv2.imshow('jiaodian', cv2.resize(rectifyImageRR, (960, 540)))
cv2.waitKey()
cv2.imwrite("jiaodian11L.png", rectifyImageRR)

# 计算相邻角点之间的距离，先计算横向相邻角点之间距离，再计算纵向相邻角点之间距离。
for i in range(0, len(cornersLO)):
    cv2.circle(grayImageL, (int(cornersLO[i][0][0]), int(cornersLO[i][0][1])), 1, (0, 0, 255))

i = 0
w_distance = []
wO_distance = []
for i in range(cornersL.shape[0]):
    if (i + 1) % 15 == 0:
        continue
    l0 = np.array([cornersL[i][0][0], cornersL[i][0][1]])
    r0 = np.array([cornersR[i][0][0], cornersR[i][0][1]])

    s0 = np.array(cv2.triangulatePoints(Pl, Pr, l0, r0)).T
    w0 = s0[0][:-1] / np.max(s0[0][-1])

    l1 = np.array([cornersL[i + 1][0][0], cornersL[i + 1][0][1]])
    r1 = np.array([cornersR[i + 1][0][0], cornersR[i + 1][0][1]])

    s1 = np.array(cv2.triangulatePoints(Pl, Pr, l1, r1)).T
    w1 = s1[0][:-1] / np.max(s1[0][-1])

    distance = np.sqrt((w0[0] - w1[0]) ** 2 + (w0[1] - w1[1]) ** 2 + (w0[2] - w1[2]) ** 2)
    w_distance.append(distance)

    lO0 = np.array([cornersLO[i][0][0], cornersLO[i][0][1]])
    rO0 = np.array([cornersRO[i][0][0], cornersRO[i][0][1]])

    sO0 = np.array(cv2.triangulatePoints(P0, P1, lO0, rO0)).T
    wO0 = sO0[0][:-1] / np.max(sO0[0][-1])

    lO1 = np.array([cornersLO[i + 1][0][0], cornersLO[i + 1][0][1]])
    rO1 = np.array([cornersRO[i + 1][0][0], cornersRO[i + 1][0][1]])

    sO1 = np.array(cv2.triangulatePoints(P0, P1, lO1, rO1)).T
    wO1 = sO1[0][:-1] / np.max(sO1[0][-1])

    distance = np.sqrt((wO0[0] - wO1[0]) ** 2 + (wO0[1] - wO1[1]) ** 2 + (wO0[2] - wO1[2]) ** 2)
    wO_distance.append(distance)

h_distance = []
hO_distance = []
for i in range(0, 15):
    # print(i)
    for j in range(i, len(cornersR), 15):
        if j >= 150:
            continue
        # print(j)
        l0 = np.array([cornersL[j][0][0], cornersL[j][0][1]])
        r0 = np.array([cornersR[j][0][0], cornersR[j][0][1]])

        s0 = np.array(cv2.triangulatePoints(Pl, Pr, l0, r0)).T
        w0 = s0[0][:-1] / np.max(s0[0][-1])

        l1 = np.array([cornersL[j + 15][0][0], cornersL[j + 15][0][1]])
        r1 = np.array([cornersR[j + 15][0][0], cornersR[j + 15][0][1]])

        s1 = np.array(cv2.triangulatePoints(Pl, Pr, l1, r1)).T
        w1 = s1[0][:-1] / np.max(s1[0][-1])  # 标定板角点三维坐标

        distance = np.sqrt((w0[0] - w1[0]) ** 2 + (w0[1] - w1[1]) ** 2 + (w0[2] - w1[2]) ** 2)
        h_distance.append(distance)

        lO0 = np.array([cornersLO[j][0][0], cornersLO[j][0][1]])
        rO0 = np.array([cornersRO[j][0][0], cornersRO[j][0][1]])

        sO0 = np.array(cv2.triangulatePoints(P0, P1, lO0, rO0)).T
        wO0 = sO0[0][:-1] / np.max(sO0[0][-1])

        lO1 = np.array([cornersLO[j + 15][0][0], cornersLO[j + 15][0][1]])
        rO1 = np.array([cornersRO[j + 15][0][0], cornersRO[j + 15][0][1]])

        sO1 = np.array(cv2.triangulatePoints(P0, P1, lO1, rO1)).T
        wO1 = sO1[0][:-1] / np.max(sO1[0][-1])

        distance = np.sqrt((wO0[0] - wO1[0]) ** 2 + (wO0[1] - wO1[1]) ** 2 + (wO0[2] - wO1[2]) ** 2)
        hO_distance.append(distance)

# 数据统计
distance = np.append(w_distance, h_distance)
distanceO = np.append(wO_distance, hO_distance)
PPM = (max(distance) - min(distance))
PPMO = (max(distanceO) - min(distanceO))
print('image', PPMO)
print('rectifyImage', PPM)

plt.figure(1)
plt.title('Error value of calculation')
plt.xlabel("The number of feature points (pcx)")
plt.ylabel("Corners distance (mm)")
plt.plot(distance, color='red', label=['rectifyppm', "%.3f" % float(PPM)])
# plt.plot(distanceO, color='gray', label=['Imageppm', "%.3f" % float(PPMO)])
plt.legend(loc='best')
plt.show()
